import torch.nn as nn


class Bottleneck(nn.Module):
    def __init__(self, in_channels, out_channels, s, downsample=False, kernel=3):
        super(Bottleneck, self).__init__()
        self.expansion = 4
        self.body = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1,
                      stride=s, padding=0, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True),
            nn.Conv2d(out_channels, out_channels, kernel,
                      stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True),
            nn.Conv2d(out_channels, out_channels * self.expansion,
                      1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels * self.expansion),
        )
        if downsample:
            self.residual = nn.Sequential(
                nn.Conv2d(in_channels, out_channels * self.expansion, 1, stride=s,
                          padding=0, bias=False),
                nn.BatchNorm2d(out_channels * self.expansion)
            )
        else:
            self.residual = None
        self.relu = nn.ReLU(True)

    def forward(self, x):
        x_body = self.body(x)
        if self.residual:
            x = self.residual(x)
        x = x_body + x
        x = self.relu(x)
        return x


class Resnet50(nn.Module):
    def __init__(self, num_classes):
        super(Resnet50, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(3, 64, 7, stride=2, padding=3, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.MaxPool2d(3, 2, padding=1),
            self.make_layer(64, 64, 3, first_block=True),
            self.make_layer(256, 128, 4),
            self.make_layer(512, 256, 6),
            self.make_layer(1024, 512, 3),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(2048, num_classes)

    def forward(self, x):
        x = self.layer(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

    def make_layer(self, in_channels, out_channels, num_block, first_block=False):
        res_block = []
        for i in range(num_block):
            if i == 0:
                if first_block:
                    res_block.append(Bottleneck(
                        in_channels, out_channels, downsample=True, s=1))
                else:
                    res_block.append(Bottleneck(
                        in_channels, out_channels, downsample=True, s=2))
            else:
                res_block.append(Bottleneck(
                    out_channels * 4, out_channels, s=1))
        return nn.Sequential(*res_block)
